# built-in
import time

# 3rd-party

from pyverbs.enums  import IBV_QPT_UD   # type: ignore
from pyverbs.enums  import IBV_WR_SEND  # type: ignore
from pyverbs.enums  import IBV_ACCESS_REMOTE_WRITE        # type: ignore
from pyverbs.enums  import IBV_ACCESS_REMOTE_READ         # type: ignore
from pyverbs.enums  import IBV_ACCESS_LOCAL_WRITE         # type: ignore

from pyverbs.cq     import CQ                             # type: ignore
from pyverbs.mr     import MR                             # type: ignore
from pyverbs.pd     import PD                             # type: ignore
from pyverbs.device import Context                        # type: ignore
from pyverbs.wr     import SGE, RecvWR, SendWR            # type: ignore
from pyverbs.qp     import QP, QPCap, QPInitAttr, QPAttr  # type: ignore
from pyverbs.addr   import AH, AHAttr, GlobalRoute        # type: ignore


# 1st-party
from utils          import parser
from conn           import SocketCM, RDMACM


RECV_WR      = 1       # 接收工作队列大小
SEND_WR      = 2       # 发送工作队列大小
GRH_LENGTH   = 40

ACCESS_FLAGS = IBV_ACCESS_LOCAL_WRITE  | \
               IBV_ACCESS_REMOTE_WRITE | \
               IBV_ACCESS_REMOTE_READ


"""
# server: 
python main.py --use_cm --ib_dev mlx5_0            \
    --mtu 4 -p 8888 -s 2147483648 -n 4 -o WRITE

    
# client: 
python main.py 100.97.0.6 --use_cm --ib_dev mlx5_0 \
    --mtu 4 -p 8888 -s 2147483648 -n 4 -o WRITE


# server
ib_send_bw -d rocep0s6 -x 1 -p 18515
ib_send_bw -d mlx5_0

# client
ib_send_bw -d rocep0s6 -x 1 192.168.1.100 -p 18515
ib_send_bw -d rocep0s6 100.97.0.6
"""

def read_mr(args: dict, mr: MR, is_server: bool) -> str:
    if args['qp_type'] == IBV_QPT_UD and is_server:
        # 从offset=GRH_LENGTH开始读取mr.length - GRH_LENGTH个字节
        return mr.read(mr.length - GRH_LENGTH, GRH_LENGTH).decode()
    else:
        # 从offset=0开始读取mr.length个(全部)字节
        return mr.read(mr.length, 0).decode()




if __name__ == "__main__":
    args = parser.parse_args()

    # 如果没有传递'server_ip'参数表示以服务端角色运行
    is_server = not bool(args["server_ip"])


    #################################################
    #! 1. Client端和Server端建立连接，初次握手
    if args['use_cm']:
        conn = RDMACM(port = args['port'], 
                    ip  = args['server_ip']
                )
    else:
        conn = SocketCM(port = args['port'], 
                        ip   = args['server_ip']
                    )

    print('-' * 80)
    print(' ' * 25, "Python PyVerbs Benchmark for RDMA")
    if is_server:
        print("以服务端身份运行...")
    else:
        print("以客户端身份运行...")
    print('-' * 80)

    if args['qp_type'] == IBV_QPT_UD and args['ope_type'] != IBV_WR_SEND:
        print("UD类型QP不支持RDMA操作...")
        conn.close()
    conn.handshake()      # 初次握手



    #################################################
    #! 2. 创建基础资源, 再次握手交换元数据
    context = Context(name=args['ib_dev'])  # 创建设备上下文
    pd      = PD(context)                   # 创建内存保护域PD
    cq      = CQ(context, 100)              # 创建完成队列CQ

    # 设置QP的参数
    cap     = QPCap(max_send_wr     = args['tx_depth'],
                    max_recv_wr     = args['rx_depth'],
                    max_send_sge    = args['sg_depth'],
                    max_recv_sge    = args['sg_depth'],
                    max_inline_data = args['inline_size'],
                )

    # 定义QP的初始化属性
    qp_init_attr = QPInitAttr(
                    qp_type = args['qp_type'], # 默认RC
                    scq = cq,
                    rcq = cq,         # 发送和接收CQ相同
                    cap = cap,
                    sq_sig_all = True # 所有WR都需要生成CQE
                )

    qp = QP(pd, qp_init_attr)

    gid = context.query_gid(port_num=1, 
                    index=args['gid_index']
                )
    remote_info = conn.handshake(gid=gid, qpn=qp.qp_num)

    #################################################
    #! 3. 创建路由并将QP设置为就绪状态
    # GlobalRoute: 配置和管理全局路由信息, 为 UD（不可靠数据报）
    #   QP（Queue Pair） 或 多播组 指定 目标端口的全局路由信息
    #     - dgid          目标地址gid
    #     - flow_label    IPv6流标签，用于数据包分类
    #     - sgid_index    源gid索引
    #     - hop_limit     跳数限制, 默认255
    #     - traffic_class 流量类别（8 位），用于 QoS 优先级（默认 0）
    gr = GlobalRoute(dgid=remote_info['gid'], 
                    sgid_index=args['gid_index']
                )
    
    # AHAttr: 地址句柄属性创建(用于配置地址句柄AH)
    # AHAttr 的主要作用是为RDMA操作（如Send/Recv）提供 目标节点的寻址和路由信息，具体包括：
    #   - 目标端口标识：通过LID（Local Identifier）或GID（Global Identifier）定位远程节点
    #   - 路径路由：指定数据包的传输路径（包括跨子网时的全局路由）
    #   - 多播支持：为多播组通信配置组地址和路由
    #   - 服务级别（SL）和路径MTU：优化数据传输的QoS和分片策略。
    # 包含如下参数:
    #
    #   - dlid   目的LID， u16， RoCE网络无效
    #   - sl     服务级别, u8, 无符号整数 
    #   - src_path_bits   src_path_bits 是一个与 LMC（LID Mask Count） 相关的参数，
    #                     用于在多路径路由（LID掩码）场景下动态调整数据包的 源LID（Local 
    #                     Identifier）。
    #           （1）LID（Local Identifier）
    #                     - LID是InfiniBand子网中唯一标识端口的16位地址（类似IP地址）。
    #                     - 传统情况下，每个端口有一个 基LID（Base LID）（例如 0x0001）。
    #           （2）LMC（LID Mask Count）
    #                     - LMC允许一个物理端口 模拟多个逻辑路径，通过扩展LID范围实现多路径路由。
    #                     - 例如：若 LMC=2，则该端口的有效LID范围为 基LID 到
    #                       基LID + (2^LMC - 1)（即 基LID 到 基LID + 3）。

    #     src_path_bits 的作用
    #            - 功能：当端口启用LMC时，数据包的实际源LID由以下公式生成：
    #                 实际源LID = 基LID | src_path_bits
    #
    #                 -  | 表示按位或（OR）操作。
    #                 - src_path_bits 是一个 8位无符号整数（0-255），用于选择逻辑路径。
    #     目的：
    #       - 负载均衡：通过不同 src_path_bits 分散流量到多条逻辑路径。
    #       - 容错：某条路径故障时，可切换到其他路径。
    #   - static_rate u8, 限制包被发送到子网的速率
    #   - is_global       否启用全局路由（GID）。若为True，需配合grh参数使用。
    #   - port_num        数据发送的物理端口号
    #   - gr              全局路由GlobalRoute
    ah_attr = AHAttr(gr=gr, is_global=1, port_num=1)

    if args['qp_type'] == IBV_QPT_UD:    # UD类型
        ah = AH(pd, attr=ah_attr)   # 创建地址句柄
        # QP的状态:
        #   - IBV_QPS_RESET    # QP 的初始状态，所有队列（SQ/RQ）被清空，无法发送或接收数据
        #   - IBV_QPS_INIT     # QP 已绑定基本属性（如端口号、QP 类型），但未建立连接, 可提交接收WR, 但不能提交发送WR
        #   - IBV_QPS_RTR      # Ready To Receive， 已准备好接收数据，但不允许发送
        #   - IBV_QPS_RTS      # Ready To Send，可同时接收和发送数据，进入正常工作状态
        #   - IBV_QPS_SQD      # 发送队列排空, 不再接受新的发送WR
        #   - IBV_QPS_SQE      # 发送队列因错误停止（如重试超限），需应用程序干预。
        #   - IBV_QPS_ERR      # QP 因严重错误（如硬件故障）不可用
        #   - IBV_QPS_UNKNOWN  # QP 状态无法确定（通常因查询失败或驱动问题） 
        qp.to_rts(QPAttr())    # 设置QP状态为RTS，可以开始发送与接收
    else:                               # 其它类型(RC, UC)
        qp_attr = QPAttr()
        qp_attr.ah_attr            = ah_attr             # 地址句柄属性
        qp_attr.dest_qp_num        = remote_info['qpn']  # 远程QP的编号
        qp_attr.path_mtu           = args['mtu']
        qp_attr.max_rd_atomic      = 1                   # 本地未完成RDMA读操作的最大数量（RC模式）。
        qp_attr.max_dest_rd_atomic = 1                   # 目标端未完成RDMA读操作的最大数量（RC模式）。
        qp_attr.qp_access_flags    = ACCESS_FLAGS

        qp.to_rtr(qp_attr) if is_server else qp.to_rts(qp_attr)
        

    conn.handshake()

    #################################################
    #! 4. 创建MR与SGL，并交换MR地址与秘钥
    role = "[Server]: " if is_server else "[Client]: "
    mr_size: int = args['mr_size']
    if is_server:
        # 在RDMA的 UD（Unreliable Datagram，不可靠数据报）模式 下，当接收的数据包
        # 包含 全局路由头（GRH，Global Route Header） 时，接收缓冲区需要 预留额外
        # 的40字节空间 来存储GRH信息。这是由UD模式的特性和InfiniBand/RoCE协议规范
        # 决定的。GRH固定为40字节大小。
        if args['qp_type'] == IBV_QPT_UD:
            mr_size += GRH_LENGTH
        content = 's'  * mr_size
    else:
        content = 'c' * mr_size

    mr_size_gb = mr_size / 1e9
    mr_size_mb = mr_size / 1e6


    mr  = MR(pd, mr_size, ACCESS_FLAGS)
    sgl = [SGE(mr.buf, mr.length, mr.lkey)]

    if args['op_type'] != IBV_WR_SEND:  # WRITE, READ
        remote_info = conn.handshake(addr=mr.buf, rkey=mr.rkey)


    #################################################
    #! 5. Client端和Server端建立连接，初次握手
    iter_num = args['iters']
    # print(f"{role} MR Content before RDMA: {content}")


    for i in range(iter_num):
        print(f"RDMA Benchmark迭代回合: {i+1}/{iter_num}", end="")

        mr.write(content, len(content))  # 将数据发送到MR缓冲区
        if is_server and args['op_type'] == IBV_WR_SEND:
            wr = RecvWR(RECV_WR, len(sgl), sgl)
            qp.post_recv(wr)
        conn.handshake()

        # Client向Server端进行RDMA操作
        if not is_server:
            wr = SendWR(SEND_WR,    # 创建WR
                    opcode  = args['op_type'], 
                    num_sge = 1, 
                    sg      = sgl
                )
            if args['qp_type'] == IBV_QPT_UD:
                wr.set_wr_ud(ah, remote_info['qpn'], 0)
            elif args['op_type'] != IBV_WR_SEND:
                wr.set_wr_rdma(remote_info['rkey'], 
                               remote_info['addr'])
            st = time.perf_counter()
            qp.post_send(wr)

            # 等待该发送事件完成
            wc_cnt, _ = cq.poll(1)
            while wc_cnt < 1:
                wc_cnt, _ = cq.poll(1)
            
            ed = time.perf_counter()

            rdma_ratio = mr_size_gb / (ed - st)
            print(f", 传输数据量{mr_size_mb:.2f}MB, "
                f"传输时间{(ed - st)* 1000:.4f}ms, "
                f"传输速度为{rdma_ratio:.2f}GB/s, wc_cnt: {wc_cnt}")
                    
        conn.handshake()

        if args['op_type'] == IBV_WR_SEND and is_server:
            wc_num, wc_list = cq.poll()
        
        if is_server:
            print(", MR Content after RDMA: " + read_mr(args, mr, is_server)[-20:])
        
    conn.handshake()   # 通信结束握手再见
    conn.close()

    print('-' * 80)